/*
 * 内存管理测试实现
 * 测试物理内存分配器、页表管理和虚拟内存功能
 */

#include "mm_test.h"
#include "kalloc.h"
#include "vm.h"
#include "riscv.h"
#include "memlayout.h"
#include "../fs/printf.h"

// 测试统计
int tests_passed = 0;
int tests_failed = 0;

/*
 * 测试物理内存分配器
 */
void test_physical_memory(void)
{
    TEST_START("Physical Memory Allocator");
    
    // 测试基本分配和释放
    test_kalloc_basic();
    
    // 测试压力情况
    test_kalloc_stress();
    
    TEST_END("Physical Memory Allocator");
}

/*
 * 基本分配测试
 */
void test_kalloc_basic(void)
{
    printf("  Testing basic allocation...\n");
    
    // 测试分配
    void *page1 = alloc_page();
    ASSERT(page1 != 0, "First allocation should succeed");
    ASSERT(((uint64)page1 & 0xFFF) == 0, "Page should be aligned");
    
    void *page2 = alloc_page();
    ASSERT(page2 != 0, "Second allocation should succeed");
    ASSERT(page1 != page2, "Different allocations should return different pages");
    
    // 测试数据写入
    *(int*)page1 = 0x12345678;
    ASSERT(*(int*)page1 == 0x12345678, "Should be able to write to allocated page");
    
    *(int*)page2 = 0x87654321;
    ASSERT(*(int*)page2 == 0x87654321, "Should be able to write to second page");
    ASSERT(*(int*)page1 == 0x12345678, "First page should not be affected");
    
    // 测试释放
    free_page(page1);
    free_page(page2);
    
    // 测试重新分配
    void *page3 = alloc_page();
    ASSERT(page3 != 0, "Allocation after free should succeed");
    
    free_page(page3);
    
    printf("  Basic allocation tests passed\n");
}

/*
 * 压力测试
 */
void test_kalloc_stress(void)
{
    printf("  Testing stress allocation...\n");
    
    #define MAX_PAGES 10
    void *pages[MAX_PAGES];
    
    // 分配多个页面
    for (int i = 0; i < MAX_PAGES; i++) {
        pages[i] = alloc_page();
        if (pages[i] == 0) {
            printf("  Warning: Could only allocate %d pages\n", i);
            break;
        }
        
        // 检查页面对齐
        ASSERT(((uint64)pages[i] & 0xFFF) == 0, "All pages should be aligned");
        
        // 写入测试数据
        *(int*)pages[i] = i;
    }
    
    // 验证数据完整性
    for (int i = 0; i < MAX_PAGES && pages[i] != 0; i++) {
        ASSERT(*(int*)pages[i] == i, "Page data should be preserved");
    }
    
    // 释放所有页面
    for (int i = 0; i < MAX_PAGES && pages[i] != 0; i++) {
        free_page(pages[i]);
    }
    
    printf("  Stress allocation tests passed\n");
}

/*
 * 测试页表功能
 */
void test_pagetable(void)
{
    TEST_START("Page Table Management");
    
    test_page_mapping();
    test_address_translation();
    
    TEST_END("Page Table Management");
}

/*
 * 页面映射测试
 */
void test_page_mapping(void)
{
    printf("  Testing page mapping...\n");
    
    // 创建测试页表
    pagetable_t pt = create_pagetable();
    ASSERT(pt != 0, "Should be able to create page table");
    
    // 分配物理页面
    void *pa = alloc_page();
    ASSERT(pa != 0, "Should be able to allocate physical page");
    
    // 测试映射
    uint64 va = 0x1000000;  // 测试虚拟地址
    int result = map_pages(pt, va, PGSIZE, (uint64)pa, PTE_R | PTE_W);
    ASSERT(result == 0, "Mapping should succeed");
    
    // 测试地址转换
    uint64 translated_pa = walkaddr(pt, va);
    ASSERT(translated_pa == (uint64)pa, "Address translation should work");
    
    // 测试重复映射（期望失败）
    result = map_pages(pt, va, PGSIZE, (uint64)pa, PTE_R | PTE_W);
    ASSERT(result != 0, "Duplicate mapping should fail");
    
    // 清理
    destroy_pagetable(pt);
    free_page(pa);
    
    printf("  Page mapping tests passed\n");
}

/*
 * 地址转换测试
 */
void test_address_translation(void)
{
    printf("  Testing address translation...\n");
    
    pagetable_t pt = create_pagetable();
    ASSERT(pt != 0, "Should be able to create page table");
    
    // 测试未映射地址
    uint64 unmapped_va = 0x2000000;
    uint64 pa = walkaddr(pt, unmapped_va);
    ASSERT(pa == 0, "Unmapped address should return 0");
    
    // 映射地址并测试
    void *phys_page = alloc_page();
    ASSERT(phys_page != 0, "Should allocate physical page");
    
    uint64 test_va = 0x3000000;
    int result = map_pages(pt, test_va, PGSIZE, (uint64)phys_page, PTE_R | PTE_W | PTE_X);
    ASSERT(result == 0, "Mapping should succeed");
    
    pa = walkaddr(pt, test_va);
    ASSERT(pa == (uint64)phys_page, "Translation should return correct physical address");
    
    // 测试页内偏移
    pa = walkaddr(pt, test_va + 0x123);
    ASSERT(pa == (uint64)phys_page, "Translation should work for page-aligned address");
    
    // 清理
    destroy_pagetable(pt);
    free_page(phys_page);
    
    printf("  Address translation tests passed\n");
}

/*
 * 测试虚拟内存激活
 */
void test_virtual_memory(void)
{
    TEST_START("Virtual Memory Activation");
    
    printf("  Before enabling paging...\n");
    
    // 检查当前SATP寄存器
    uint64 satp_before = r_satp();
    printf("  SATP before: 0x%lx\n", satp_before);
    
    // 初始化内核虚拟内存
    kvminit();
    ASSERT(kernel_pagetable != 0, "Kernel page table should be created");
    
    printf("  Kernel page table created at: 0x%lx\n", (uint64)kernel_pagetable);
    
    // 激活内核页表
    kvminithart();
    
    uint64 satp_after = r_satp();
    printf("  SATP after: 0x%lx\n", satp_after);
    ASSERT(satp_after != satp_before, "SATP should change after activation");
    
    printf("  After enabling paging...\n");
    
    // 测试内核代码仍然可执行
    printf("  Kernel code still executable\n");
    
    // 测试内核数据仍然可访问
    static int test_data = 42;
    ASSERT(test_data == 42, "Kernel data should still be accessible");
    
    // 测试堆栈仍然可用
    int stack_var = 123;
    ASSERT(stack_var == 123, "Stack should still be accessible");
    
    printf("  Virtual memory activation successful\n");
    
    TEST_END("Virtual Memory Activation");
}

/*
 * 运行所有内存管理测试
 */
void run_all_mm_tests(void)
{
    printf("\n=== Memory Management Tests ===\n");
    
    tests_passed = 0;
    tests_failed = 0;
    
    // 首先初始化物理内存分配器
    printf("Initializing physical memory allocator...\n");
    pmm_init();
    
    // 运行测试
    test_physical_memory();
    test_pagetable();
    test_virtual_memory();
    
    // 输出测试结果
    printf("\n=== Test Results ===\n");
    printf("Tests passed: %d\n", tests_passed);
    printf("Tests failed: %d\n", tests_failed);
    
    if (tests_failed == 0) {
        printf("All memory management tests PASSED!\n");
    } else {
        printf("Some tests FAILED!\n");
    }
    
    printf("=== End of Tests ===\n\n");
}